package com.gitee.wsl.ext.list

/**
 * Returns a new list where each element is preceded by its results in [sourceToTarget]. The first
 * element will return no values in [sourceToTarget].
 *
 * @param sourceToTarget a function that returns nodes that should precede the argument in the result.
 *
 * Implementation inspiration: https://www.interviewcake.com/concept/java/topological-sort
 */
fun <T> List<T>.topologicalSort(sourceToTarget: (T) -> Iterable<T>): List<T> {
    // Build a reverse index, from targets to sources.
    val targetToSources = mutableMapOf<T, MutableSet<T>>()
    val queue = ArrayDeque<T>()
    for (source in this) {
        var hasTargets = false
        for (target in sourceToTarget(source)) {
            val set = targetToSources.getOrPut(target) { mutableSetOf() }
            set += source
            hasTargets = true
        }

        // No targets means all this source's targets are satisfied, queue it up.
        if (!hasTargets) {
            queue += source
        }
    }

    val result = mutableListOf<T>()
    while (queue.isNotEmpty()) {
        val node = queue.removeFirst()
        result += node

        val potentiallySatisfied = targetToSources[node] ?: setOf()
        for (source in potentiallySatisfied) {
            // If all a source's targets are satisfied, queue up the source.
            if (source !in queue &&
                sourceToTarget(source).all { target -> target in result || target in queue }
            ) {
                queue += source
            }
        }
    }

    require(result.size == this.size) {
        "No topological ordering is possible for $this"
    }

    return result
}

fun <T> List<T>.isTopologicallySorted(sourceToTarget: (T) -> Iterable<T>): Boolean {
    val seenNodes = mutableSetOf<T>()
    for (node in this) {
        if (sourceToTarget(node).any { it !in seenNodes }) return false
        seenNodes.add(node)
    }
    return true
}

/**
 * Extension function to perform quick sort on a list of comparable elements.
 *
 * @receiver List<T> The list to be sorted.
 * @return List<T> A new list containing the sorted elements.
 */
fun <T : Comparable<T>> List<T>.quickSort(): List<T> =
    when {
        size < 2 -> this
        else -> {
            val pivot = first()
            val (smaller, greater) = drop(1).partition { it <= pivot }
            smaller.quickSort() + pivot + greater.quickSort()
        }
    }

fun <T> List<T>.isSortedWith(comparator: Comparator<T>): Boolean =
    asSequence().zipWithNext { a, b -> comparator.compare(a, b) <= 0 }.all { it }
